// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <melon/cportability.h>

#include <pollux/expression/function_signature.h>
#include <pollux/functions/lib/checked_arithmetic_lmpl.h>
#include <pollux/functions/lib/aggregates/decimal_aggregate.h>
#include <pollux/functions/lib/aggregates/simple_numeric_aggregate.h>

namespace kumo::pollux::functions::aggregate {
    template<
        typename TInput,
        typename TAccumulator,
        typename ResultType,
        bool Overflow>
    class SumAggregateBase
            : public SimpleNumericAggregate<TInput, TAccumulator, ResultType> {
        using BaseAggregate =
        SimpleNumericAggregate<TInput, TAccumulator, ResultType>;

    public:
        explicit SumAggregateBase(TypePtr resultType) : BaseAggregate(resultType) {
        }

        int32_t accumulatorFixedWidthSize() const override {
            return sizeof(TAccumulator);
        }

        int32_t accumulatorAlignmentSize() const override {
            return 1;
        }

        void extractValues(char **groups, int32_t numGroups, VectorPtr *result)
        override {
            BaseAggregate::template doExtractValues<ResultType>(
                groups, numGroups, result, [&](char *group) {
                    // 'ResultType' and 'TAccumulator' might not be same such as sum(real)
                    // and we do an explicit type conversion here.
                    return (ResultType) (*BaseAggregate::Aggregate::template value<
                        TAccumulator>(group));
                });
        }

        void extractAccumulators(char **groups, int32_t numGroups, VectorPtr *result)
        override {
            BaseAggregate::template doExtractValues<TAccumulator>(
                groups, numGroups, result, [&](char *group) {
                    return *BaseAggregate::Aggregate::template value<TAccumulator>(group);
                });
        }

        void addRawInput(
            char **groups,
            const SelectivityVector &rows,
            const std::vector<VectorPtr> &args,
            bool mayPushdown) override {
            updateInternal<TAccumulator>(groups, rows, args, mayPushdown);
        }

        void addIntermediateResults(
            char **groups,
            const SelectivityVector &rows,
            const std::vector<VectorPtr> &args,
            bool mayPushdown) override {
            updateInternal<TAccumulator, TAccumulator>(groups, rows, args, mayPushdown);
        }

        void addSingleGroupRawInput(
            char *group,
            const SelectivityVector &rows,
            const std::vector<VectorPtr> &args,
            bool mayPushdown) override {
            BaseAggregate::template updateOneGroup<TAccumulator>(
                group,
                rows,
                args[0],
                &updateSingleValue<TAccumulator>,
                &updateDuplicateValues<TAccumulator>,
                mayPushdown,
                TAccumulator(0));
        }

        void addSingleGroupIntermediateResults(
            char *group,
            const SelectivityVector &rows,
            const std::vector<VectorPtr> &args,
            bool mayPushdown) override {
            BaseAggregate::template updateOneGroup<TAccumulator, TAccumulator>(
                group,
                rows,
                args[0],
                &updateSingleValue<TAccumulator>,
                &updateDuplicateValues<TAccumulator>,
                mayPushdown,
                TAccumulator(0));
        }

    protected:
        // TData is used to store the updated sum state. It can be either
        // TAccumulator or TResult, which in most cases are the same, but for
        // sum(real) can differ. TValue is used to decode the sum input 'args'.
        // It can be either TAccumulator or TInput, which is most cases are the same
        // but for sum(real) can differ.
        template<typename TData, typename TValue = TInput>
        void updateInternal(
            char **groups,
            const SelectivityVector &rows,
            const std::vector<VectorPtr> &args,
            bool mayPushdown) {
            const auto &arg = args[0];

            if (mayPushdown && arg->is_lazy()) {
                BaseAggregate::template pushdown<
                    kumo::pollux::aggregate::SumHook<TData, Overflow> >(
                    groups, rows, arg);
                return;
            }

            if (exec::Aggregate::numNulls_) {
                BaseAggregate::template updateGroups<true, TData, TValue>(
                    groups, rows, arg, &updateSingleValue<TData>, false);
            } else {
                BaseAggregate::template updateGroups<false, TData, TValue>(
                    groups, rows, arg, &updateSingleValue<TData>, false);
            }
        }

        void initializeNewGroupsInternal(
            char **groups,
            melon::Range<const vector_size_t *> indices) override {
            exec::Aggregate::setAllNulls(groups, indices);
            for (auto i: indices) {
                *exec::Aggregate::value<TAccumulator>(groups[i]) = 0;
            }
        }

    private:
        template<typename TData>
        static void updateSingleValue(TData &result, TData value) {
            pollux::aggregate::SumHook<TData, Overflow>::add(result, value);
        }

        // Disable undefined behavior sanitizer to not fail on signed integer
        // overflow.
        template<typename TData>
#if defined(MELON_DISABLE_UNDEFINED_BEHAVIOR_SANITIZER)
        MELON_DISABLE_UNDEFINED_BEHAVIOR_SANITIZER("signed-integer-overflow")
#endif
        static void updateDuplicateValues(TData &result, TData value, int n) {
            if constexpr (
                (std::is_same_v<TData, int64_t> && Overflow) ||
                std::is_same_v<TData, double> || std::is_same_v<TData, float>) {
                result += n * value;
            } else {
                result = functions::checkedPlus<TData>(
                    result, functions::checkedMultiply<TData>(TData(n), value));
            }
        }
    };

    template<typename TInputType>
    class DecimalSumAggregate
            : public functions::aggregate::DecimalAggregate<int128_t, TInputType> {
    public:
        explicit DecimalSumAggregate(TypePtr resultType)
            : functions::aggregate::DecimalAggregate<int128_t, TInputType>(
                resultType) {
        }

        virtual int128_t computeFinalValue(
            functions::aggregate::LongDecimalWithOverflowState *accumulator) final {
            auto sum = DecimalUtil::adjustSumForOverflow(
                accumulator->sum, accumulator->overflow);
            POLLUX_USER_CHECK(sum.has_value(), "Decimal overflow");
            DecimalUtil::valueInRange(sum.value());
            return sum.value();
        }
    };
} // namespace kumo::pollux::functions::aggregate
